// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <boost/numeric/conversion/cast.hpp>
#include <chrono>
#include <optional>
#include <pollux/common/base/doubles.h>
#include <pollux/external/date/date.h>
#include <pollux/functions/lib/date_time_formatter.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>
#include <pollux/type/timestamp.h>
#include <pollux/type/timestamp_conversion.h>
#include <pollux/type/tz/time_zone_map.h>

namespace kumo::pollux::functions {
    namespace {
        constexpr double kNanosecondsInSecond = 1'000'000'000;
        constexpr int64_t kNanosecondsInMillisecond = 1'000'000;
        constexpr int64_t kMillisecondsInSecond = 1'000;
    } // namespace

    MELON_ALWAYS_INLINE double toUnixtime(const Timestamp &timestamp) {
        double result = timestamp.getSeconds();
        result += static_cast<double>(timestamp.getNanos()) / kNanosecondsInSecond;
        return result;
    }

    MELON_ALWAYS_INLINE Timestamp fromUnixtime(double unixtime) {
        if (MELON_UNLIKELY(std::isnan(unixtime))) {
            return Timestamp(0, 0);
        }

        static const int64_t kMin = std::numeric_limits<int64_t>::min();

        if (MELON_UNLIKELY(unixtime >= kMinDoubleAboveInt64Max)) {
            return Timestamp::maxMillis();
        }

        if (MELON_UNLIKELY(unixtime <= kMin)) {
            return Timestamp::minMillis();
        }

        if (MELON_UNLIKELY(std::isinf(unixtime))) {
            return unixtime < 0 ? Timestamp::minMillis() : Timestamp::maxMillis();
        }

        auto seconds = std::floor(unixtime);
        auto milliseconds = std::llround((unixtime - seconds) * kMillisInSecond);
        if (MELON_UNLIKELY(milliseconds == kMillisInSecond)) {
            ++seconds;
            milliseconds = 0;
        }
        return Timestamp(seconds, milliseconds * kNanosecondsInMillisecond);
    }

    MELON_ALWAYS_INLINE boost::int64_t fromUnixtime(
        double unixtime,
        int16_t timeZoneId) {
        if (MELON_UNLIKELY(std::isnan(unixtime))) {
            return pack(0, timeZoneId);
        }

        static const int64_t kMin = std::numeric_limits<int64_t>::min();

        if (MELON_UNLIKELY(unixtime >= kMinDoubleAboveInt64Max)) {
            return pack(std::numeric_limits<int64_t>::max(), timeZoneId);
        }

        if (MELON_UNLIKELY(unixtime <= kMin)) {
            return pack(std::numeric_limits<int64_t>::min(), timeZoneId);
        }

        if (MELON_UNLIKELY(std::isinf(unixtime))) {
            return unixtime < 0
                       ? pack(std::numeric_limits<int64_t>::min(), timeZoneId)
                       : pack(std::numeric_limits<int64_t>::max(), timeZoneId);
        }

        return pack(std::llround(unixtime * kMillisecondsInSecond), timeZoneId);
    }

    // Year, quarter or month are not uniformly incremented. Months have different
    // total days, and leap years have more days than the rest. If the new year,
    // quarter or month has less total days than the given one, it will be coerced
    // to use the valid last day of the new month. This could result in weird
    // arithmetic behavior. For example,
    //
    // 2022-01-30 + (1 month) = 2022-02-28
    // 2022-02-28 - (1 month) = 2022-01-28
    //
    // 2022-08-31 + (1 quarter) = 2022-11-30
    // 2022-11-30 - (1 quarter) = 2022-08-30
    //
    // 2020-02-29 + (1 year) = 2021-02-28
    // 2021-02-28 - (1 year) = 2020-02-28
    MELON_ALWAYS_INLINE
    int32_t
    addToDate(const int32_t input, const DateTimeUnit unit, const int32_t value) {
        // TODO(gaoge): Handle overflow and underflow with 64-bit representation
        if (value == 0) {
            return input;
        }

        const std::chrono::time_point<std::chrono::system_clock, date::days> inDate{
            date::days(input)
        };
        std::chrono::time_point<std::chrono::system_clock, date::days> outDate;

        if (unit == DateTimeUnit::kDay) {
            outDate = inDate + date::days(value);
        } else if (unit == DateTimeUnit::kWeek) {
            outDate = inDate + date::days(value * 7);
        } else {
            const date::year_month_day inCalDate(inDate);
            date::year_month_day outCalDate;

            if (unit == DateTimeUnit::kMonth) {
                outCalDate = inCalDate + date::months(value);
            } else if (unit == DateTimeUnit::kQuarter) {
                outCalDate = inCalDate + date::months(3 * value);
            } else if (unit == DateTimeUnit::kYear) {
                outCalDate = inCalDate + date::years(value);
            } else {
                POLLUX_UNREACHABLE();
            }

            if (!outCalDate.ok()) {
                outCalDate = outCalDate.year() / outCalDate.month() / date::last;
            }
            outDate = date::sys_days{outCalDate};
        }

        return outDate.time_since_epoch().count();
    }

    MELON_ALWAYS_INLINE Timestamp addToTimestamp(
        const Timestamp &timestamp,
        const DateTimeUnit unit,
        const int32_t value) {
        // TODO(gaoge): Handle overflow and underflow with 64-bit representation
        if (value == 0) {
            return timestamp;
        }

        const std::chrono::
                time_point<std::chrono::system_clock, std::chrono::milliseconds>
                inTimestamp(std::chrono::milliseconds(timestamp.toMillis()));
        std::chrono::time_point<std::chrono::system_clock, std::chrono::milliseconds>
                outTimestamp;

        switch (unit) {
            // Year, quarter or month are not uniformly incremented in terms of number
            // of days. So we treat them differently.
            case DateTimeUnit::kYear:
            case DateTimeUnit::kQuarter:
            case DateTimeUnit::kMonth:
            case DateTimeUnit::kDay: {
                const int32_t inDate =
                        std::chrono::duration_cast<date::days>(inTimestamp.time_since_epoch())
                        .count();
                const int32_t outDate = addToDate(inDate, unit, value);

                outTimestamp = inTimestamp + date::days(outDate - inDate);
                break;
            }
            case DateTimeUnit::kHour: {
                outTimestamp = inTimestamp + std::chrono::hours(value);
                break;
            }
            case DateTimeUnit::kMinute: {
                outTimestamp = inTimestamp + std::chrono::minutes(value);
                break;
            }
            case DateTimeUnit::kSecond: {
                outTimestamp = inTimestamp + std::chrono::seconds(value);
                break;
            }
            case DateTimeUnit::kMillisecond: {
                outTimestamp = inTimestamp + std::chrono::milliseconds(value);
                break;
            }
            case DateTimeUnit::kWeek: {
                const int32_t inDate =
                        std::chrono::duration_cast<date::days>(inTimestamp.time_since_epoch())
                        .count();
                const int32_t outDate = addToDate(inDate, DateTimeUnit::kDay, 7 * value);

                outTimestamp = inTimestamp + date::days(outDate - inDate);
                break;
            }
            default:
                POLLUX_UNREACHABLE("Unsupported datetime unit");
        }

        Timestamp milliTimestamp =
                Timestamp::fromMillis(outTimestamp.time_since_epoch().count());
        return Timestamp(
            milliTimestamp.getSeconds(),
            milliTimestamp.getNanos() +
            timestamp.getNanos() % kNanosecondsInMillisecond);
    }

    // If time zone is provided, use it for the arithmetic operation (convert to it,
    // apply operation, then convert back to UTC).
    MELON_ALWAYS_INLINE Timestamp addToTimestamp(
        const Timestamp &timestamp,
        const DateTimeUnit unit,
        const int32_t value,
        const tz::TimeZone *timeZone) {
        if (timeZone == nullptr) {
            return addToTimestamp(timestamp, unit, value);
        } else {
            Timestamp zonedTimestamp = timestamp;
            zonedTimestamp.toTimezone(*timeZone);
            auto resultTimestamp = addToTimestamp(zonedTimestamp, unit, value);
            resultTimestamp.toGMT(*timeZone);
            return resultTimestamp;
        }
    }

    MELON_ALWAYS_INLINE int64_t addToTimestampWithTimezone(
        int64_t timestampWithTimezone,
        const DateTimeUnit unit,
        const int32_t value) {
        {
            int64_t finalSysMs;
            if (unit < DateTimeUnit::kDay) {
                auto originalTimestamp = unpackTimestampUtc(timestampWithTimezone);
                finalSysMs =
                        addToTimestamp(originalTimestamp, unit, (int32_t) value).toMillis();
            } else {
                // Use local time to handle crossing daylight savings time boundaries.
                // E.g. the "day" when the clock moves back an hour is 25 hours long, and
                // the day it moves forward is 23 hours long. Daylight savings time
                // doesn't affect time units less than a day, and will produce incorrect
                // results if we use local time.
                const tz::TimeZone *timeZone =
                        tz::locateZone(unpackZoneKeyId(timestampWithTimezone));
                auto originalTimestamp =
                        Timestamp::fromMillis(timeZone
                            ->to_local(std::chrono::milliseconds(
                                unpackMillisUtc(timestampWithTimezone)))
                            .count());
                auto updatedTimeStamp =
                        addToTimestamp(originalTimestamp, unit, (int32_t) value);
                updatedTimeStamp = Timestamp(
                    timeZone
                    ->correct_nonexistent_time(
                        std::chrono::seconds(updatedTimeStamp.getSeconds()))
                    .count(),
                    updatedTimeStamp.getNanos());
                finalSysMs =
                        timeZone
                        ->to_sys(
                            std::chrono::milliseconds(updatedTimeStamp.toMillis()),
                            tz::TimeZone::TChoose::kEarliest)
                        .count();
            }

            return pack(finalSysMs, unpackZoneKeyId(timestampWithTimezone));
        }
    }

    MELON_ALWAYS_INLINE int64_t diffTimestamp(
        const DateTimeUnit unit,
        const Timestamp &fromTimestamp,
        const Timestamp &toTimestamp) {
        // TODO(gaoge): Handle overflow and underflow with 64-bit representation
        if (fromTimestamp == toTimestamp) {
            return 0;
        }

        const int8_t sign = fromTimestamp < toTimestamp ? 1 : -1;

        // fromTimepoint is less than or equal to toTimepoint
        const std::chrono::
                time_point<std::chrono::system_clock, std::chrono::milliseconds>
                fromTimepoint(std::chrono::milliseconds(
                    std::min(fromTimestamp, toTimestamp).toMillis()));
        const std::chrono::
                time_point<std::chrono::system_clock, std::chrono::milliseconds>
                toTimepoint(std::chrono::milliseconds(
                    std::max(fromTimestamp, toTimestamp).toMillis()));

        // Millisecond, second, minute, hour and day have fixed conversion ratio
        switch (unit) {
            case DateTimeUnit::kMillisecond: {
                return sign *
                       std::chrono::duration_cast<std::chrono::milliseconds>(
                           toTimepoint - fromTimepoint)
                       .count();
            }
            case DateTimeUnit::kSecond: {
                return sign *
                       std::chrono::duration_cast<std::chrono::seconds>(
                           toTimepoint - fromTimepoint)
                       .count();
            }
            case DateTimeUnit::kMinute: {
                return sign *
                       std::chrono::duration_cast<std::chrono::minutes>(
                           toTimepoint - fromTimepoint)
                       .count();
            }
            case DateTimeUnit::kHour: {
                return sign *
                       std::chrono::duration_cast<std::chrono::hours>(
                           toTimepoint - fromTimepoint)
                       .count();
            }
            case DateTimeUnit::kDay: {
                return sign *
                       std::chrono::duration_cast<date::days>(toTimepoint - fromTimepoint)
                       .count();
            }
            case DateTimeUnit::kWeek: {
                return sign *
                       std::chrono::duration_cast<date::days>(toTimepoint - fromTimepoint)
                       .count() /
                       7;
            }
            default:
                break;
        }

        // Month, quarter and year do not have fixed conversion ratio. Ex. a month can
        // have 28, 29, 30 or 31 days. A year can have 365 or 366 days.
        const std::chrono::time_point<std::chrono::system_clock, date::days>
                fromDaysTimepoint = std::chrono::floor<date::days>(fromTimepoint);
        const std::chrono::time_point<std::chrono::system_clock, date::days>
                toDaysTimepoint = std::chrono::floor<date::days>(toTimepoint);
        const date::year_month_day fromCalDate(fromDaysTimepoint);
        const date::year_month_day toCalDate(toDaysTimepoint);
        const uint64_t fromTimeInstantOfDay =
                std::chrono::duration_cast<std::chrono::milliseconds>(
                    fromTimepoint - fromDaysTimepoint)
                .count();

        uint64_t toTimeInstantOfDay = 0;
        uint64_t toTimePointMillis = toTimepoint.time_since_epoch().count();
        uint64_t toDaysTimepointMillis =
                std::chrono::
                time_point<std::chrono::system_clock, std::chrono::milliseconds>(
                    toDaysTimepoint)
                .time_since_epoch()
                .count();
        bool overflow = __builtin_sub_overflow(
            toTimePointMillis, toDaysTimepointMillis, &toTimeInstantOfDay);
        POLLUX_USER_CHECK_EQ(
            overflow,
            false,
            "{} - {} Causes arithmetic overflow: {} - {}",
            fromTimestamp.toString(),
            toTimestamp.toString(),
            toTimePointMillis,
            toDaysTimepointMillis);
        const uint8_t fromDay = static_cast<unsigned>(fromCalDate.day()),
                fromMonth = static_cast<unsigned>(fromCalDate.month());
        const uint8_t toDay = static_cast<unsigned>(toCalDate.day()),
                toMonth = static_cast<unsigned>(toCalDate.month());
        const date::year_month_day toCalLastYearMonthDay(
            toCalDate.year() / toCalDate.month() / date::last);
        const uint8_t toLastYearMonthDay =
                static_cast<unsigned>(toCalLastYearMonthDay.day());

        if (unit == DateTimeUnit::kMonth || unit == DateTimeUnit::kQuarter) {
            int64_t diff =
                    (int64_t(toCalDate.year()) - int64_t(fromCalDate.year())) * 12 +
                    int(toMonth) - int(fromMonth);

            if ((toDay != toLastYearMonthDay && fromDay > toDay) ||
                (fromDay == toDay && fromTimeInstantOfDay > toTimeInstantOfDay)) {
                diff--;
            }

            diff = (unit == DateTimeUnit::kMonth) ? diff : diff / 3;
            return sign * diff;
        }

        if (unit == DateTimeUnit::kYear) {
            int64_t diff = (toCalDate.year() - fromCalDate.year()).count();

            if (fromMonth > toMonth ||
                (fromMonth == toMonth && fromDay > toDay &&
                 toDay != toLastYearMonthDay) ||
                (fromMonth == toMonth && fromDay == toDay &&
                 fromTimeInstantOfDay > toTimeInstantOfDay)) {
                diff--;
            }
            return sign * diff;
        }

        POLLUX_UNREACHABLE("Unsupported datetime unit");
    }

    MELON_ALWAYS_INLINE int64_t diffTimestampWithTimeZone(
        const DateTimeUnit unit,
        const int64_t fromTimestampWithTimeZone,
        const int64_t toTimestampWithTimeZone) {
        auto fromTimeZoneId = unpackZoneKeyId(fromTimestampWithTimeZone);
        auto toTimeZoneId = unpackZoneKeyId(toTimestampWithTimeZone);
        POLLUX_CHECK_EQ(
            fromTimeZoneId,
            toTimeZoneId,
            "diffTimestampWithTimeZone must receive timestamps in the same time zone.");

        Timestamp fromTimestamp;
        Timestamp toTimestamp;

        if (unit < DateTimeUnit::kDay) {
            fromTimestamp = unpackTimestampUtc(fromTimestampWithTimeZone);
            toTimestamp = unpackTimestampUtc(toTimestampWithTimeZone);
        } else {
            // Use local time to handle crossing daylight savings time boundaries.
            // E.g. the "day" when the clock moves back an hour is 25 hours long, and
            // the day it moves forward is 23 hours long. Daylight savings time
            // doesn't affect time units less than a day, and will produce incorrect
            // results if we use local time.
            const tz::TimeZone *timeZone = tz::locateZone(fromTimeZoneId);
            fromTimestamp = Timestamp::fromMillis(
                timeZone
                ->to_local(std::chrono::milliseconds(
                    unpackMillisUtc(fromTimestampWithTimeZone)))
                .count());
            toTimestamp =
                    Timestamp::fromMillis(timeZone
                        ->to_local(std::chrono::milliseconds(
                            unpackMillisUtc(toTimestampWithTimeZone)))
                        .count());
        }

        return diffTimestamp(unit, fromTimestamp, toTimestamp);
    }

    MELON_ALWAYS_INLINE
    int64_t diffDate(
        const DateTimeUnit unit,
        const int32_t fromDate,
        const int32_t toDate) {
        if (fromDate == toDate) {
            return 0;
        }
        return diffTimestamp(
            unit,
            // prevent overflow
            Timestamp((int64_t) fromDate * util::kSecsPerDay, 0),
            Timestamp((int64_t) toDate * util::kSecsPerDay, 0));
    }

    MELON_ALWAYS_INLINE int64_t
    valueOfTimeUnitToMillis(const double value, std::string_view unit) {
        double convertedValue = value;
        if (unit == "ns") {
            convertedValue = convertedValue * std::milli::den / std::nano::den;
        } else if (unit == "us") {
            convertedValue = convertedValue * std::milli::den / std::micro::den;
        } else if (unit == "ms") {
        } else if (unit == "s") {
            convertedValue = convertedValue * std::milli::den;
        } else if (unit == "m") {
            convertedValue = convertedValue * 60 * std::milli::den;
        } else if (unit == "h") {
            convertedValue = convertedValue * 3600 * std::milli::den;
        } else if (unit == "d") {
            convertedValue = convertedValue * 86400 * std::milli::den;
        } else {
            POLLUX_USER_FAIL("Unknown time unit: {}", unit);
        }
        try {
            return boost::numeric_cast<int64_t>(std::round(convertedValue));
        } catch (const boost::bad_numeric_cast &) {
            POLLUX_USER_FAIL(
                "Value in {} unit is too large to be represented in ms unit as an int64_t",
                unit,
                value);
        }
        POLLUX_UNREACHABLE();
    }
} // namespace kumo::pollux::functions
